import numpy as np
import time

# 球队实力先验分布 (0-3分别代表四种实力等级)
PA = np.array([0.3, 0.3, 0.2, 0.2])  # 队伍A的实力分布
PB = np.array([0.4, 0.4, 0.1, 0.1])  # 队伍B的实力分布
PC = np.array([0.2, 0.2, 0.3, 0.3])  # 队伍C的实力分布

# 比赛结果条件概率表 PS[主队实力][客队实力] = [胜, 平, 负]
PS = np.array([
    [[0.2, 0.6, 0.2], [0.1, 0.3, 0.6], [0.05, 0.2, 0.75], [0.01, 0.1, 0.89]],
    [[0.6, 0.3, 0.1], [0.2, 0.6, 0.2], [0.1, 0.3, 0.6], [0.05, 0.2, 0.75]],
    [[0.75, 0.2, 0.05], [0.6, 0.3, 0.1], [0.2, 0.6, 0.2], [0.1, 0.3, 0.6]],
    [[0.89, 0.1, 0.01], [0.75, 0.2, 0.05], [0.6, 0.3, 0.1], [0.2, 0.6, 0.2]]
])


def normalize(x):
    """归一化概率分布"""
    x = np.array(x, dtype=float)
    return x / np.sum(x)


def exact_inference():
    """
    精确推理方法
    时间复杂度: O(4^3 * 3) = O(192)
    """
    res = np.zeros(3)
    # 遍历所有可能的实力组合
    for XA in range(4):
        for XB in range(4):
            for XC in range(4):
                # 计算联合概率 P(XA,XB,XC,sAB=0,sAC=1,sBC)
                joint_prob = PA[XA] * PB[XB] * PC[XC] * \
                             PS[XA][XB][0] * \
                             PS[XA][XC][1]  # P(sAC=1|XA,XC)

                # 累加不同sBC的概率
                for sBC in range(3):
                    res[sBC] += joint_prob * PS[XB][XC][sBC]  # P(sBC|XB,XC)

    return normalize(res)


def rejection_sampling(n=5000):
    """
    拒绝采样方法
    时间复杂度: 取决于接受率，最坏情况O(∞)
    """
    res = np.zeros(3)
    accepted = 0
    attempts = 0

    while accepted < n:
        attempts += 1
        # 从先验分布采样实力
        XA = np.random.choice(4, p=PA)
        XB = np.random.choice(4, p=PB)
        XC = np.random.choice(4, p=PC)

        # 采样比赛结果
        sAB = np.random.choice(3, p=PS[XA][XB])
        sAC = np.random.choice(3, p=PS[XA][XC])

        # 只接受符合证据的样本(sAB=0, sAC=1)
        if sAB == 0 and sAC == 1:
            sBC = np.random.choice(3, p=PS[XB][XC])
            res[sBC] += 1
            accepted += 1

    acceptance_rate = accepted / attempts
    return normalize(res), acceptance_rate


def likelihood_weighting(n=5000):
    """
    似然加权采样方法
    时间复杂度: O(n)
    """
    res = np.zeros(3)
    total_weight = 0.0

    for _ in range(n):
        weight = 1.0
        # 采样实力(非证据变量)
        XA = np.random.choice(4, p=PA)
        XB = np.random.choice(4, p=PB)
        XC = np.random.choice(4, p=PC)

        # 固定证据变量并计算权重
        weight *= PS[XA][XB][0]  # sAB=0的似然
        weight *= PS[XA][XC][1]  # sAC=1的似然

        # 采样查询变量
        sBC = np.random.choice(3, p=PS[XB][XC])
        res[sBC] += weight
        total_weight += weight

    # 计算有效样本量
    if total_weight > 0:
        normalized_weights = res / total_weight
        effective_samples = 1.0 / np.sum(normalized_weights ** 2)
    else:
        effective_samples = 0

    return normalize(res), effective_samples


def gibbs_sampling(n=5000, burn_in=1000):
    """
    Gibbs采样方法
    时间复杂度: O((burn_in + n) * 3)
    """
    res = np.zeros(3)

    # 初始化变量
    XA, XB, XC = np.random.choice(4), np.random.choice(4), np.random.choice(4)
    sBC = np.random.choice(3)  # 初始sBC值

    for i in range(burn_in + n):
        # 1. 采样XA (条件依赖: XB, XC, sAB=0, sAC=1)
        prob_XA = [PA[a] * PS[a][XB][0] * PS[a][XC][1] for a in range(4)]
        XA = np.random.choice(4, p=normalize(prob_XA))

        # 2. 采样XB (条件依赖: XA, XC, sAB=0, sBC)
        prob_XB = [PB[b] * PS[XA][b][0] * PS[b][XC][sBC] for b in range(4)]
        XB = np.random.choice(4, p=normalize(prob_XB))

        # 3. 采样XC (条件依赖: XA, XB, sAC=1, sBC)
        prob_XC = [PC[c] * PS[XA][c][1] * PS[XB][c][sBC] for c in range(4)]
        XC = np.random.choice(4, p=normalize(prob_XC))

        # 4. 采样sBC
        sBC = np.random.choice(3, p=PS[XB][XC])

        # 只保留burn-in期后的样本
        if i >= burn_in:
            res[sBC] += 1

    return normalize(res)


def analyze_performance():
    """性能分析函数"""
    print("=== 开始性能分析 ===")

    # 计算精确解作为基准
    start = time.time()
    exact_result = exact_inference()
    exact_time = time.time() - start
    print(f"精确解: {exact_result} (耗时: {exact_time:.4f}s)\n")

    # 测试拒绝采样
    start = time.time()
    rej_result, rej_rate = rejection_sampling(2000)
    rej_time = time.time() - start
    rej_error = np.sum(np.abs(rej_result - exact_result))
    print(f"拒绝采样结果: {rej_result}")
    print(f"接受率: {rej_rate:.2%} 耗时: {rej_time:.4f}s L1误差: {rej_error:.6f}\n")

    # 测试似然加权
    start = time.time()
    lw_result, eff_samples = likelihood_weighting(2000)
    lw_time = time.time() - start
    lw_error = np.sum(np.abs(lw_result - exact_result))
    print(f"似然加权结果: {lw_result}")
    print(f"有效样本量: {eff_samples:.1f} 耗时: {lw_time:.4f}s L1误差: {lw_error:.6f}\n")

    # 测试Gibbs采样
    start = time.time()
    gibbs_result = gibbs_sampling(2000)
    gibbs_time = time.time() - start
    gibbs_error = np.sum(np.abs(gibbs_result - exact_result))
    print(f"Gibbs采样结果: {gibbs_result}")
    print(f"耗时: {gibbs_time:.4f}s L1误差: {gibbs_error:.6f}")


if __name__ == '__main__':
    analyze_performance()
